import { create } from 'zustand';
import { persist } from 'zustand/middleware';
import dayjs, { Dayjs } from 'dayjs';
import isBetween from 'dayjs/plugin/isBetween';
import { type GPU, type Server, type Reservation, type ReservationFormData, type GPUStatusInfo } from '../types';

dayjs.extend(isBetween);

interface GPUState {
  servers: Server[];
  reservations: Reservation[];
  selectedGpu: GPU | null;
  isModalVisible: boolean;
  
  // Actions
  setSelectedGpu: (gpu: GPU | null) => void;
  setModalVisible: (visible: boolean) => void;
  addReservation: (reservation: Omit<Reservation, 'id'>) => void;
  deleteReservation: (id: number) => void;
  checkTimeConflict: (gpuId: number, startTime: Dayjs, endTime: Dayjs) => boolean;
  createReservationFromForm: (formData: ReservationFormData) => Promise<void>;
  getGpuReservations: (gpuId: number) => Reservation[];
  getGpuStatus: (gpuId: number) => GPUStatusInfo;
  addServer: (server: Omit<Server, 'id' | 'gpus'>) => void;
  addGpuToServer: (serverId: number, gpu: Omit<GPU, 'id'>) => void;
}

// 初始服务器数据
const initialServers: Server[] = [
  {
    id: 1,
    name: "AI训练服务器-01",
    location: "机房A-机架12",
    gpus: Array.from({ length: 8 }, (_, i) => ({
      id: i + 1,
      name: `GPU-${String(i + 1).padStart(2, '0')}`,
      model: "NVIDIA A100",
      memory: "80GB",
      serverId: 1
    }))
  },
  {
    id: 2,
    name: "数据分析服务器-02",
    location: "机房B-机架05",
    gpus: Array.from({ length: 8 }, (_, i) => ({
      id: i + 9,
      name: `GPU-${String(i + 9).padStart(2, '0')}`,
      model: "NVIDIA V100",
      memory: "32GB",
      serverId: 2
    }))
  }
];

export const useGPUStore = create<GPUState>()(
  persist(
    (set, get) => ({
      servers: initialServers,
      reservations: [],
      selectedGpu: null,
      isModalVisible: false,

      setSelectedGpu: (gpu) => set({ selectedGpu: gpu }),

      setModalVisible: (visible) => set({ isModalVisible: visible }),

      addReservation: (reservationData) => {
        const newReservation: Reservation = {
          ...reservationData,
          id: Date.now()
        };
        
        set((state) => ({
          reservations: [...state.reservations, newReservation]
        }));
      },

      deleteReservation: (id) => {
        set((state) => ({
          reservations: state.reservations.filter(r => r.id !== id)
        }));
      },

      checkTimeConflict: (gpuId, startTime, endTime) => {
        const { reservations } = get();
        
        return reservations.some(reservation => {
          if (reservation.gpuId !== gpuId) return false;
          
          return (
            startTime.isBetween(reservation.startTime, reservation.endTime, null, '[]') ||
            endTime.isBetween(reservation.startTime, reservation.endTime, null, '[]') ||
            (startTime.isBefore(reservation.startTime) && endTime.isAfter(reservation.endTime))
          );
        });
      },

      createReservationFromForm: async (formData) => {
        const { selectedGpu, checkTimeConflict, addReservation } = get();
        
        if (!selectedGpu) {
          throw new Error('请先选择一张GPU卡');
        }

        const { applicant, purpose, contact, project, dateRange, timeRange } = formData;
        const [startDate, endDate] = dateRange;
        const [startTime, endTime] = timeRange;

        const startDateTime = startDate
          .hour(startTime.hour())
          .minute(startTime.minute())
          .second(0);
        
        const endDateTime = endDate
          .hour(endTime.hour())
          .minute(endTime.minute())
          .second(0);

        if (startDateTime.isAfter(endDateTime)) {
          throw new Error('开始时间不能晚于结束时间');
        }

        if (startDateTime.isBefore(dayjs())) {
          throw new Error('不能预定过去的时间');
        }

        if (checkTimeConflict(selectedGpu.id, startDateTime, endDateTime)) {
          throw new Error('该时间段已被预定，请选择其他时间');
        }

        addReservation({
          gpuId: selectedGpu.id,
          applicant,
          purpose,
          contact,
          project,
          startTime: startDateTime,
          endTime: endDateTime
        });
      },

      getGpuReservations: (gpuId) => {
        const { reservations } = get();
        return reservations.filter(r => r.gpuId === gpuId);
      },

      getGpuStatus: (gpuId) => {
        const { reservations } = get();
        const now = dayjs();
        const gpuReservations = reservations.filter(r => r.gpuId === gpuId);
        
        const currentReservation = gpuReservations.find(r => 
          now.isBetween(r.startTime, r.endTime, null, '[]')
        );

        if (currentReservation) {
          return { 
            status: 'in-use', 
            user: currentReservation.applicant,
            reservation: currentReservation
          };
        }

        const upcomingReservation = gpuReservations.find(r => 
          r.startTime.isAfter(now)
        );

        if (upcomingReservation) {
          return { 
            status: 'reserved', 
            user: upcomingReservation.applicant,
            reservation: upcomingReservation
          };
        }

        // 5%的几率显示为维护中
        if (Math.random() < 0.05) {
          return { status: 'maintenance' };
        }

        return { status: 'available' };
      },

      addServer: (serverData) => {
        const newServer: Server = {
          ...serverData,
          id: Date.now(),
          gpus: []
        };
        
        set((state) => ({
          servers: [...state.servers, newServer]
        }));
      },

      addGpuToServer: (serverId, gpuData) => {
        const newGpu: GPU = {
          ...gpuData,
          id: Date.now(),
          serverId
        };
        
        set((state) => ({
          servers: state.servers.map(server => 
            server.id === serverId 
              ? { ...server, gpus: [...server.gpus, newGpu] }
              : server
          )
        }));
      }
    }),
    {
      name: 'gpu-storage',
      partialize: (state) => ({ 
        reservations: state.reservations,
        servers: state.servers
      }),
      onRehydrateStorage: () => (state) => {
        if (state) {
          // 将存储的字符串时间转换回Dayjs对象
          state.reservations = state.reservations.map((r: any) => ({
            ...r,
            startTime: dayjs(r.startTime),
            endTime: dayjs(r.endTime)
          }));
        }
      },
    }
  )
);